library(optparse)
library(tidyverse)
library(magrittr)
library(RJSONIO)
library(extrafont)
library(circlize)

configGet <- function(configData, key) {
  configData %>%
    filter(arg == key) %>%
    .$value
}

option_list <- list(
  make_option("--tablepath", default = "AllMet_Raw.txt", type = "character", help = "Data matrix table path"),
  make_option("--imagepackagepath", default = "./", type = "character", help = "Output the folder package in which the image resides"),
  make_option("--input_nrow", default = "all", type = "character", help = "the number of rows entered"),
  make_option("--input_ncol", default = "all", type = "character", help = "the number of cols entered"),
  make_option("--treecolfile", default = " ", type = "character", help = "col tree file path"),
  make_option("--labfilepath", default = " ", type = "character", help = "Custom grid label file"),
  make_option("--checked_lab_names", default = "TRUE", type = "logical", help = "wheather to checked file col and row names"),
  make_option("--lab_number_color", default = "black", type = "character", help = "color of lab"),
  make_option("--showxyc_names", default = "TRUE:TRUE:FALSE", type = "character", help = "Whether to display the XY axis and data text"),
  make_option("--show_annotataion_names", default = "TRUE:TRUE", type = "character", help = "Whether to display col row names"),
  make_option("--imageformt", default = "pdf", type = "character", help = "pdf,tiff,png"),
  make_option("--imageName", default = "heatmap", type = "character", help = "picture name"),
  make_option("--pc", default = "plot_config.txt", type = "character", help = "config file"),
  make_option("--cc", default = "calculate_config.txt", type = "character", help = "config file")
)
opt <- parse_args(OptionParser(option_list = option_list))

args <- commandArgs(trailingOnly = F)
scriptPath <- dirname(sub("--file=", "", args[grep("--file", args)]))
source(str_c(scriptPath, "/parse_config_json.R"))

suppressMessages(library("pheatmap"))
suppressMessages(library("ggplot2"))
suppressMessages(library("dplyr"))
suppressMessages(library("ggtree"))
suppressMessages(library("ape"))
suppressMessages(library("ggplotify"))

fenge <- function(str) {
  str <- strsplit(str, ":")[[1]]
  return(str)
}

fenge <- function(str, sp = ":", fix = FALSE) {
  return(strsplit(str, sp, fix)[[1]])
}

fenge2 <- function(str) {
  str <- strsplit(str, ",")[[1]]
  return(str)
}

fenge3 <- function(str) {
  str <- strsplit(str, "-")[[1]]
  return(str)
}

topTriangle <- function(data) {
  ll <- min(ncol(data), nrow(data))
  data <- as.data.frame(data); data[, ncol(data) + 1] <- NA; data[nrow(data) + 1,] <- NA
  for (k in 1:ll) {
    data[(k + 1):nrow(data), k] <- NA
  }
  data <- data[1:(nrow(data) - 1), 1:(ncol(data) - 1)]
  return(as.matrix(data))
}

bottomTriangle <- function(data) {
  ll <- min(ncol(data), nrow(data))
  for (k in 2:ll) {
    data[1:(k - 1), k] <- NA
  }

  return(data)
}

left <- function(data) {
  mid <- round(((ncol(data) / 2) + 0.1))
  data[, (mid + 1):ncol(data)] <- NA
  return(data)
}

right <- function(data) {
  mid <- round(((ncol(data) / 2) + 0.1))
  data[, 1:(mid - 1)] <- NA
  return(data)
}

top <- function(data) {
  mid <- round(((nrow(data) / 2) + 0.1))
  data[(mid + 1):nrow(data),] <- NA
  return(data)
}

bottom <- function(data) {
  mid <- round(((nrow(data) / 2) + 0.1))
  data[1:(mid - 1),] <- NA
  return(data)
}

changetype <- function(data, str) {
  if (str == "left") { data <- left(data) }
  if (str == "right") { data <- right(data) }
  if (str == "top") { data <- top(data) }
  if (str == "bottom") { data <- bottom(data) }
  if (str == "topTriangle") { data <- topTriangle(data) }
  if (str == "bottomTriangle") { data <- bottomTriangle(data) }
  return(data)
}

hebinglist <- function(d1, d2) {
  if (is.na(d1)) { return(d2) }
  if (is.na(d2)) { return(d1) }
  if (!is.na(d2) && !is.na(d1)) {
    result <- list()
    names <- c(names(d1), names(d2))

    for (i in seq_along(d1)) {
      result[[i]] <- d1[[i]]
    }
    ll <- length(result)
    for (i in seq_along(d2)) {
      result[[ll + i]] <- d2[[i]]
    }
    names(result) <- names
    return(result)
  }
}

refreshAnnoRowColors <- function(list, pcFile) {
  jsonStr <- RJSONIO::toJSON(list, collapse = "")
  plotConfigJson <- fromJSON(pcFile)
  plotConfigJson$color$annoRowColors <- jsonStr
  write_json(plotConfigJson, pcFile, auto_unbox = TRUE)
}

refreshAnnoColColors <- function(list, pcFile) {
  jsonStr <- RJSONIO::toJSON(list, collapse = "")
  plotConfigJson <- fromJSON(pcFile)
  plotConfigJson$color$annoColColors <- jsonStr
  write_json(plotConfigJson, pcFile, auto_unbox = TRUE)
}

gg_color_hue <- function(n) {
  hues = seq(15, 375, length = n + 1)
  hcl(h = hues, l = 65, c = 100)[1:n]
}

getClassCols <- function(n, rev = F) {
  inCols <- c("#00FFFF", "#0099FF", "#8A4198", "#FF6F00", "#CC9999", "#FFFF00", "#00FF7F", "#CCFFCC", "#336666",
              "#FF6666", "#FFCC33", "#CC3399")
  extraColors <- c("#CD0000", "#3A89CC", "#769C30", "#D99536", "#7B0078", "#BFBC3B", "#6E8B3D", "#00688B", "#C10077", "#CAAA76",
                   "#EEEE00", "#458B00", "#8B4513", "#008B8B", "#6E8B3D", "#8B7D6B", "#7FFF00", "#CDBA96", "#ADFF2F")
  orignalCols <- if (rev) {
    c(rev(inCols), rev(extraColors))
  } else c(inCols, extraColors)
  on <- length(orignalCols)
  if (n <= on) {
    orignalCols[1:n]
  }else {
    ggCols <- gg_color_hue(n - on)
    c(orignalCols, ggCols)
  }
}

data <- read.delim(opt$tablepath, row.names = 1, check.names = FALSE)
data1 <- data

showPart <- plotConfigJson$layout$showPart
rowClusterMethod <- plotConfigJson$layout$rowClusterMethod
rowCluster <- plotConfigJson$layout$rowCluster %>%
  as.logical()
splitNum <- plotConfigJson$layout$splitNum
colCluster <- plotConfigJson$layout$colCluster != "F"
colClusterMethod <- plotConfigJson$layout$colClusterMethod
colSplitNum <- plotConfigJson$layout$colSplitNum
convMethod <- plotConfigJson$layout$convMethod
normal <- plotConfigJson$layout$normal
log <- plotConfigJson$layout$log
hasBorder <- plotConfigJson$layout$hasBorder
width <- plotConfigJson$layout$width %>%
  as.numeric()
height <- plotConfigJson$layout$height %>%
  as.numeric()

gradientNum <- plotConfigJson$color$gradientNum %>%
  as.numeric()
naColor <- plotConfigJson$color$naColor
borderColor <- plotConfigJson$color$borderColor
finalBorderColor <- if (hasBorder == "none") {
  NA
}else borderColor
showRowName <- plotConfigJson$font$showRowName
showColName <- plotConfigJson$font$showColName
showNumber <- plotConfigJson$font$showNumber
xRotate <- plotConfigJson$font$xRotate
xFont <- plotConfigJson$font$xFont %>%
  as.numeric()
yFont <- plotConfigJson$font$yFont %>%
  as.numeric()
rowTreeHeight <- plotConfigJson$layout$rowTreeHeight
colTreeHeight <- plotConfigJson$layout$colTreeHeight
cellFont <- plotConfigJson$font$cellFont
mainTitle <- plotConfigJson$font$mainTitle
annoRowColors <- plotConfigJson$color$annoRowColors
annoColColors <- plotConfigJson$color$annoColColors
fontFamily <- plotConfigJson$font$fontFamily
baseFamily <- fontFamily

calculateConfigData <- read_tsv(opt$cc, col_names = F, col_types = "cc") %>%
  set_colnames(c("arg", "value"))
hasRowGroup <- configGet(calculateConfigData, "hasRowGroup") %>%
  as.logical()
rowGroupFile <- if (hasRowGroup) "row_group.txt" else ""
hasColGroup <- configGet(calculateConfigData, "hasColGroup") %>%
  as.logical()
colGroupFile <- if (hasColGroup) "col_group.txt" else ""

if (opt$input_nrow != "all") {
  inputrow <- fenge2(opt$input_nrow)
  ilength <- length(inputrow)
  nr <- NULL
  for (i in 1:ilength) {
    if (length(strsplit(inputrow[i], "-")[[1]]) == 1) {
      z <- as.numeric(strsplit(inputrow[i], "-")[[1]])
      nr[length(nr) + 1] <- z
    }else {
      z <- as.numeric(fenge3(inputrow[i]))
      nr[length(nr) + 1:(length(nr) + 1 + z[2] - z[1])] <- z[1]:z[2]
    }
  }
  nr <- unique(nr)
  data <- data[nr,]
}
if (opt$input_ncol != "all") {
  input_ncol <- fenge2(opt$input_ncol)
  ilength <- length(input_ncol)
  nr <- NULL
  for (i in 1:ilength) {
    if (length(strsplit(input_ncol[i], "-")[[1]]) == 1) {
      z <- as.numeric(strsplit(input_ncol[i], "-")[[1]])
      nr[length(nr) + 1] <- z
    }else {
      z <- as.numeric(fenge3(input_ncol[i]))
      nr[length(nr) + 1:(length(nr) + 1 + z[2] - z[1])] <- z[1]:z[2]
    }
  }
  nr <- unique(nr)
  data <- data[, nr] }

if (convMethod == "log") {
  if (log == "lg2") { data <- log2(data + 1) }
  if (log == "lg10") { data <- log10(data + 1) }
}else if (convMethod == "zScale") {
  if (normal == "row") { data <- t(scale(t(data))); data[is.nan(data)] <- NA }
  if (normal == "col") { data <- scale(data); data[is.nan(data)] <- NA }
}

if (opt$treecolfile == " ") {
  cluster_cols <- as.logical(colCluster)
  if (as.logical(colCluster)) {
    cluster_cols <- hclust(dist(t(data)), method = colClusterMethod)
    cluster_cols <- as.phylo.hclust(cluster_cols)
    write.tree(file = paste0(opt$imagepackagepath, "/col.tre"), cluster_cols)
    cluster_cols <- as.hclust(read.tree(paste0(opt$imagepackagepath, "/col.tre")))
    hh <- ggtree(cluster_cols)$data %>% arrange(y)
    hh <- unique(hh$label)
    hh <- as.character(na.omit(hh))
    ktt <- match(hh, colnames(data))
    data <- data[, ktt]
  }
}else {
  cluster_cols <- as.hclust(read.tree(opt$treecolfile))
  hh <- ggtree(cluster_cols)$data %>% arrange(y)
  hh <- unique(hh$label)
  hh <- as.character(na.omit(hh))
  ktt <- match(hh, colnames(data))
  data <- data[, ktt]
}

if (rowCluster) {
  meta_dist <- dist(as.matrix(1 - (data)), method = "euclidean")
  hc_meta <- hclust(meta_dist, method = rowClusterMethod)
}else {
  hc_meta <- FALSE
}

annotation1 <- NA
print(rowGroupFile)
if (rowGroupFile != "") {
  annotation1 <- read.delim(rowGroupFile, row.names = 1, check.names = FALSE)
}
annotation2 <- NA
if (colGroupFile != "") {
  annotation2 <- read_tsv(colGroupFile) %>%
    mutate_at(vars(-c("Sample")), ~as.character(.x)) %>%
    column_to_rownames("Sample")
}

if (convMethod == "zScale" & normal == "row") {
  color <- colorRampPalette(finalColors)(7)
  medianNum <- ceiling(gradientNum / 2)
  min <- data %>%
    discard(~is.infinite(.x)) %>%
    min()
  print(min)
  max <- max(data)
  minColorsF <- colorRamp2(seq(min, 0, length.out = 4), color[1:4])
  maxColorsF <- colorRamp2(seq(0, max, length.out = 4), color[4:7])
  minVs = seq(min, 0, length.out = medianNum)
  minColors <- minColorsF(minVs)
  maxVs = seq(0, max, length.out = medianNum)
  maxColors <- maxColorsF(maxVs)
  colors <- c(minColors, maxColors) %>%
    unique()
  breaks <- c(minVs, maxVs) %>%
    unique()
}else {
  colors <- colorRampPalette(finalColors)(gradientNum)
  breaks <- NA
}

showxyc_names <- fenge(opt$showxyc_names)
main <- mainTitle
if (as.logical(showRowName)) { labels_row <- NULL }else { labels_row <- rep("", nrow(data)) }
if (as.logical(showColName)) { labels_col <- NULL }else { labels_col <- rep("", ncol(data)) }
show_annotataion_names <- fenge(opt$show_annotataion_names)
display <- as.logical(showNumber)
if (opt$labfilepath != " ") {
  display <- read.delim(opt$labfilepath, row.names = 1, check.names = FALSE)
  if (as.logical(opt$checked_lab_names)) {
    datac <- rownames(data)
    displayc <- rownames(display)
    display <- display[match(datac, displayc),]
    datac <- colnames(data)
    displayc <- colnames(display)
    display <- display[, match(datac, displayc)]
  }
}
type <- showPart
if (class(display) != "logical") {
  display <- changetype(display, showPart)
  display[is.na(display)] <- ""
}
if (showPart != "full") {
  data <- changetype(data, showPart)
}

annotation_colors <- list()
pcFile <- opt$pc
anno1IsNA <- is.na(annotation1)
anno2IsNA <- is.na(annotation2)
if (!anno1IsNA || !anno2IsNA) {
  annoRowColorList <- if (annoRowColors=="") {
    if (anno1IsNA) {
      NA
    }else {
      columnNames <- colnames(annotation1)
      list <- columnNames %>%
        purrr::set_names() %>%
        map(function(columnName) {
          classes <- annotation1[, columnName] %>%
            unique()
          cols <- getClassCols(length(classes)) %>%
            purrr::set_names(., classes)
          cols
        })
      refreshAnnoRowColors(list, pcFile)
      list
    }
  }else {
    RJSONIO::fromJSON(annoRowColors)
  }
  annoColColorList <- if (annoColColors=="") {
    if (anno2IsNA) {
      NA
    }else {
      columnNames <- colnames(annotation2)
      list <- columnNames %>%
        purrr::set_names() %>%
        map(function(columnName) {
          classes <- annotation2[, columnName] %>%
            unique()
          cols <- getClassCols(length(classes), rev = T) %>%
            purrr::set_names(., classes)
          cols
        })
      refreshAnnoColColors(list, pcFile)
      list
    }
  }else {
    RJSONIO::fromJSON(annoColColors)
  }
  annoColorList <- hebinglist(annoRowColorList, annoColColorList)
  annotation_colors <- annoColorList
}else {
  annotation_colors <- NA
}
print(annotation_colors)
pp <- pheatmap(data
  , cluster_rows = hc_meta
  , treeheight_row = rowTreeHeight
  , cluster_cols = cluster_cols
  , treeheight_col = colTreeHeight
  , border_color = finalBorderColor
  , display_numbers = display
  , show_colnames = as.logical(showxyc_names[1])
  , show_rownames = as.logical(showxyc_names[2])
  , color = colors, breaks = breaks,
  , fontsize_row = yFont
  , fontsize_col = xFont
  , cutree_rows = splitNum
  , cutree_cols = colSplitNum
  , fontsize_number = cellFont
  , annotation_col = annotation2
  , annotation_row = annotation1
  , number_color = opt$lab_number_color
  , main = main
  , na_col = naColor
  , labels_row = labels_row
  , labels_col = labels_col
  , annotation_names_col = as.logical(show_annotataion_names[2])
  , annotation_names_row = as.logical(show_annotataion_names[1])
  , angle_col = xRotate
  , filename = paste0(opt$imagepackagepath, "/", opt$imageName, ".", opt$imageformt)
  , width = width
  , height = height
  , annotation_colors = annotation_colors,
               fontfamily = baseFamily
)

rowNum<-nrow(data)
colNum<-ncol(data)
plotConfigJson <- fromJSON(pcFile)
plotConfigJson$layout$rowNum <- rowNum
plotConfigJson$layout$colNum <- colNum
write_json(plotConfigJson, pcFile, auto_unbox = TRUE)

